{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "342942fd-79e2-4fa7-bafe-4896d8b578a1",
   "metadata": {},
   "source": [
    "# Wikipedia Semantic Search with Cohere + Weaviate\n",
    "This is starter code that you can use to search 10 million vectors from wikipedia embedded with Cohere's multilingual model and hosted as a Weaviate public dataset. This dataset contains 1M vectors in each of the Wikipedia sites in these languages: English, German, French, Spanish, Italian, Japanese, Arabic, Chinese (Simplified), Korean, Hindi \\[respective language codes: `en, de, fr, es, it, ja, ar, zh, ko, hi`\\]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43debcc4-eb17-4bc7-8210-96a0905eec1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: upgrade to \"cohere>5\"",
"!pip install \"cohere<5\" weaviate-client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "8f840bfe-d592-414e-90ce-a04dfc3cd87e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import weaviate\n",
    "\n",
    "# Add your Cohere API key here\n",
    "cohere_api_key = ''\n",
    "\n",
    "# Connect to the Weaviate demo databse containing 10M wikipedia vectors\n",
    "# This uses a public READ-ONLY Weaviate API key\n",
    "auth_config = weaviate.auth.AuthApiKey(api_key=\"76320a90-53d8-42bc-b41d-678647c6672e\") \n",
    "client = weaviate.Client(\n",
    "    url=\"https://cohere-demo.weaviate.network/\",\n",
    "    auth_client_secret=auth_config,\n",
    "    additional_headers={\n",
    "        \"X-Cohere-Api-Key\": cohere_api_key,\n",
    "    }\n",
    ")\n",
    "\n",
    "client.is_ready() #check if True"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da70e927-3c9a-4a9f-8d8c-2778b2ab42b4",
   "metadata": {},
   "source": [
    "Let's now define the search function that queries our vector database. Optionally, we want the ability to filter by language."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "3987c132-06ca-4ee4-a1a9-c6ac04fd7f09",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def semantic_serch(query, results_lang=''):\n",
    "    \"\"\" \n",
    "    Query the vectors database and return the top results. \n",
    "\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "        query: str\n",
    "            The search query\n",
    "            \n",
    "        results_lang: str (optional)\n",
    "            Retrieve results only in the specified language.\n",
    "            The demo dataset has those languages:\n",
    "            en, de, fr, es, it, ja, ar, zh, ko, hi\n",
    "\n",
    "    \"\"\"\n",
    "    \n",
    "    nearText = {\"concepts\": [query]}\n",
    "    properties = [\"text\", \"title\", \"url\", \"views\", \"lang\", \"_additional {distance}\"]\n",
    "\n",
    "    # To filter by language\n",
    "    if results_lang != '':\n",
    "        where_filter = {\n",
    "        \"path\": [\"lang\"],\n",
    "        \"operator\": \"Equal\",\n",
    "        \"valueString\": results_lang\n",
    "        }\n",
    "        response = (\n",
    "            client.query\n",
    "            .get(\"Articles\", properties)\n",
    "            .with_where(where_filter)\n",
    "            .with_near_text(nearText)\n",
    "            .with_limit(5)\n",
    "            .do()\n",
    "        )\n",
    "        \n",
    "    # Search all languages\n",
    "    else:\n",
    "        response = (\n",
    "            client.query\n",
    "            .get(\"Articles\", properties)\n",
    "            .with_near_text(nearText)\n",
    "            .with_limit(5)\n",
    "            .do()\n",
    "        )\n",
    "\n",
    "\n",
    "    result = response['data']['Get']['Articles']\n",
    "\n",
    "    return result\n",
    "\n",
    "\n",
    "def print_result(result):\n",
    "    \"\"\" Print results with colorful formatting \"\"\"\n",
    "    for item in result:\n",
    "        print(f\"\\033[95m{item['title']} ({item['views']}) {item['_additional']['distance']}\\033[0m\")\n",
    "        print(f\"\\033[4m{item['url']}\\033[0m\")\n",
    "        print(item['text'])\n",
    "        print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b0cb81c-8e56-48cf-a32c-f050afe646e4",
   "metadata": {},
   "source": [
    "We can now query the databse with any query we want. In the background, Weaviate uses your Cohere API key to embed the query, then retrun the most relevant passages to the query."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "5102ea41-fe8f-4ced-8b23-229df185d447",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[95mThe Adam Project (3000) -147.31755\u001b[0m\n",
      "\u001b[4mhttps://en.wikipedia.org/wiki?curid=65867428\u001b[0m\n",
      "Due to a safety feature preventing him from flying because of his injuries, Adam must bring along the younger Adam and use his DNA to enter his jet. They both are soon attacked by Maya Sorian, the leader of the dystopian world, and her assistant Christos, but are saved by Laura, who had faked her death and stayed off-grid in an unknown location. After surviving the attack and comparing notes, Laura and the Adams realize that after the invention of time travel by Louis Reed and his subsequent death, Sorian had monopolized the discovery. During her visit to 2018, Laura learned Sorian frequently came and advised her past self in order to secure her future wealth and power. To protect her secret, Sorian ordered Laura's death. Although Laura survived the assassination attempt, destruction of her time jet left her stranded in the past. The sudden arrival of Sorian's goons interrupts the reunion, and Laura fights off the attack long enough for the two Adams to escape to 2018.\n",
      "\n",
      "\u001b[95mKang the Conqueror (2000) -146.57275\u001b[0m\n",
      "\u001b[4mhttps://en.wikipedia.org/wiki?curid=393437\u001b[0m\n",
      "Nathaniel Richards, a 31st-century scholar and descendant of Reed Richards' time traveling father Nathaniel, becomes fascinated with history and discovers the time travel technology created by Victor von Doom, another possible ancestor of his. He then travels back in time to ancient Egypt aboard a Sphinx-shaped timeship and becomes the Pharaoh Rama-Tut, with plans to claim En Sabah Nur—the mutant destined to become Apocalypse—as his heir. Rama-Tut's rule is cut short when he is defeated by the time-displaced Fantastic Four. An embittered Nathaniel Richards travels forward to the 20th century where he meets Doctor Doom, whom he believes might be his ancestor. He later designs an armor based on Doom's and, calling himself the Scarlet Centurion, pits the Avengers team against alternate-reality counterparts. He plans to dispose of all of them, but the Avengers manage to force him from the timeline.\n",
      "\n",
      "\u001b[95mBack to the Future (3000) -146.4269\u001b[0m\n",
      "\u001b[4mhttps://en.wikipedia.org/wiki?curid=42993\u001b[0m\n",
      "That night, Marty meets his eccentric scientist friend, Emmett \"Doc\" Brown, in the Twin Pines mall parking lot. Doc unveils a time machine built from a modified DeLorean, powered by plutonium he swindled from Libyan terrorists. After Doc inputs a destination time of November5, 1955—the day he first conceived his time travel invention—the terrorists arrive unexpectedly and gun Doc down. Marty flees in the DeLorean, inadvertently activating time travel when he reaches .\n",
      "\n",
      "\u001b[95mTime (2000) -146.41129\u001b[0m\n",
      "\u001b[4mhttps://en.wikipedia.org/wiki?curid=30012\u001b[0m\n",
      "Time travel is the concept of moving backwards or forwards to different points in time, in a manner analogous to moving through space, and different from the normal \"flow\" of time to an earthbound observer. In this view, all points in time (including future times) \"persist\" in some way. Time travel has been a plot device in fiction since the 19th century. Travelling backwards or forwards in time has never been verified as a process, and doing so presents many theoretical problems and contradictive logic which to date have not been overcome. Any technological device, whether fictional or hypothetical, that is used to achieve time travel is known as a time machine.\n",
      "\n",
      "\u001b[95mIn Time (2000) -145.93015\u001b[0m\n",
      "\u001b[4mhttps://en.wikipedia.org/wiki?curid=29446866\u001b[0m\n",
      "Will and Sylvia rob Weis' time banks, giving the time capsules to the needy. They soon realize they can't significantly change anything, as prices are raised faster to compensate for the extra time. Fortis' gang ambushes them intending to collect the reward for their capture, but Will kills Fortis and his gang. Will and Sylvia then decide to rob Weis' vault of a one-million year capsule. Leon chases them back to Dayton but fails to stop them from distributing the stolen time; Leon times out, having neglected to collect his day's salary. Will and Sylvia nearly time out themselves but survive by taking Leon's salary.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "query_result = semantic_serch(\"time travel plot twist\")\n",
    "\n",
    "# Print out the result\n",
    "print_result(query_result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0dd5114-4604-4c3c-8077-29ac3d5524f4",
   "metadata": {},
   "source": [
    "## Filtering by language\n",
    "If we're interested in results in only one language, we can specify it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "67e9be31-9754-4d74-8600-43fa0840e26c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[95m時空の旅人 (500) -144.16002\u001b[0m\n",
      "\u001b[4mhttps://ja.wikipedia.org/wiki?curid=523464\u001b[0m\n",
      "バスは1868年の攘夷戦争で娘と夫を亡くした老婆の営む茶店に降り立つ。一時は老婆は一行を盗人と間違えて襲い掛かるも、ホクベンを亡き夫だと思い込んだことで一転して歓迎する。しかしそこへジロを追うタイムマシンがあらわれ、やむなく一行はバスに乗って走り去る。追い縋る老婆を見捨てられずバスを飛び降りたホクベンだが、直後にタイムマシンに攫われてしまった。\n",
      "\n",
      "\u001b[95m親殺しのパラドックス (700) -144.11864\u001b[0m\n",
      "\u001b[4mhttps://ja.wikipedia.org/wiki?curid=71910\u001b[0m\n",
      "パラドックスを防ぐアイデアとして、時間旅行者は元々の歴史とは異なる並行宇宙に行くのだと解釈するもので、上の科学的理論で述べたのと同じ考え方であり、SFにもよく見られる。歴史改変SFにあるタイムトラベル参照。\n",
      "\n",
      "\u001b[95mタイムトラベル (1000) -143.70331\u001b[0m\n",
      "\u001b[4mhttps://ja.wikipedia.org/wiki?curid=1971274\u001b[0m\n",
      "タイムパラドックスの矛盾を説明するため、タイムトラベル者による歴史の改変で時間軸が分岐し元の世界と並行した別の世界が生まれるとするパラレルワールドの概念がある。この概念を発展させ、タイムトラベル者の介在がなくとも歴史上の重要なポイントで世界が枝分かれしていると解釈する立場もある。この概念を大幅に作品に取り入れた最初期の小説に、可能性として存在する二つの歴史「ジョンバール」と「ギロンチ」の抗争を描いた、ジャック・ウィリアムスンの『航時軍団』（The Legion of Time、1938年）がある。\n",
      "\n",
      "\u001b[95mタイムトラベル (1000) -143.69884\u001b[0m\n",
      "\u001b[4mhttps://ja.wikipedia.org/wiki?curid=1971274\u001b[0m\n",
      "タイムトラベラーが主人公であるマーク・トウェインの「アーサー王宮廷のコネチカット・ヤンキー」や、天使が未来の書物を携えて現れるサミュエル・マッデンの「20世紀回想」など、SFというカテゴリが明確なものとして育つ以前から、タイムトラベルをテーマにした物語は創られている。\n",
      "\n",
      "\u001b[95mタイムトラベル (1000) -143.61562\u001b[0m\n",
      "\u001b[4mhttps://ja.wikipedia.org/wiki?curid=1971274\u001b[0m\n",
      "タイムパラドックス（Time Paradox / 時間の逆説）は、タイムトラベルに伴う矛盾や変化のことであり、物語のテーマとしてしばしば扱われる。具体的には、タイムトラベルした過去で現代（相対的未来）に存在する事象を改変した場合、その事象における過去と現代の存在や状況、因果関係の不一致という逆説が生じることに着目したものである。\n",
      "\n"
     ]
    }
   ],
   "source": [
    "query_result = semantic_serch(\"time travel plot twist\", results_lang='ja')\n",
    "\n",
    "# Print out the result\n",
    "print_result(query_result)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
